package com.lemon.components.ant.util.sqlparser;

import java.util.ArrayList;
import java.util.List;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource;
import com.alibaba.druid.sql.ast.statement.SQLTableSource;
import com.alibaba.druid.sql.ast.statement.SQLUnionQuery;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;

public class CustomerMySqlASTVisitorAdapter extends MySqlASTVisitorAdapter {

	private final List<SQLTableInfo> tableInfos = new ArrayList<>();
    private boolean isFirstTable = true; // 用于标识是否是主表

    @Override
    public boolean visit(SQLExprTableSource x) {
        SQLExpr expr = x.getExpr();
        if (expr != null) {
            String tableName = expr.toString();
            // 第一个表是主表
            tableInfos.add(new SQLTableInfo(tableName, isFirstTable));
            isFirstTable = false; // 后续表不是主表
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLSelectQueryBlock x) {
        // 处理 FROM 子句
        SQLTableSource from = x.getFrom();
        if (from != null) {
            from.accept(this);
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLUnionQuery x) {
        // 处理联合查询
        SQLSelectQuery left = x.getLeft();
        SQLSelectQuery right = x.getRight();
        if (left != null) {
            left.accept(this);
        }
        if (right != null) {
            right.accept(this);
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLSubqueryTableSource x) {
        // 处理子查询
        SQLSelect select = x.getSelect();
        if (select != null) {
            select.accept(this);
        }
        return super.visit(x);
    }

    @Override
    public boolean visit(SQLJoinTableSource x) {
        // 处理 JOIN 子句
        SQLTableSource left = x.getLeft();
        SQLTableSource right = x.getRight();
        if (left != null) {
            left.accept(this);
        }
        if (right != null) {
            right.accept(this);
        }
        return super.visit(x);
    }

    public List<SQLTableInfo> getTableInfos() {
        return tableInfos;
    }

	public static class SQLTableInfo {

		private String tableName;
		private boolean isMainTable;

		public SQLTableInfo(String tableName, boolean isMainTable) {
			this.tableName = tableName;
			this.isMainTable = isMainTable;
		}

		public String getTableName() {
			return tableName;
		}

		public boolean isMainTable() {
			return isMainTable;
		}

		@Override
		public String toString() {
			return "TableInfo{" + "tableName='" + tableName + '\'' + ", isMainTable=" + isMainTable + '}';
		}
	}
}
